# 导入必要的库
from sklearn.model_selection import train_test_split  # 用于拆分数据集为训练集和测试集
import cv2
import pickle
import numpy as np
from sklearn.ensemble import RandomForestClassifier

# 读取砂岩截面图1及其对应分区图(灰度图)  
image = cv2.imread('D:/Sandstone_imgs/Sandstone_1.tif', cv2.IMREAD_GRAYSCALE)
print("图像形状:",image.shape[:2])  
segment = cv2.imread('D:/Sandstone_imgs/Sandstone_1_segment.tif', cv2.IMREAD_GRAYSCALE)  
  
# 获取图像数据和标签  
X = image.reshape(-1, 1)  
y = segment.ravel() 
print("完成从砂岩截面图1及其对应分区中获取X和y") 
  
# 划分训练集和测试集  
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)  
print("完成train_test_split")  

# 训练随机森林模型  
clf = RandomForestClassifier()  
clf.fit(X_train, y_train)  
print("完成随机森林模型clf的训练")
  
# 输出模型准确率  
print(f"准确率: {clf.score(X_test, y_test)}")  
  
# 保存模型到硬盘
with open('clf_model.pkl', 'wb') as f:
    pickle.dump(clf, f)
print("已保存随机森林模型clf到硬盘")